import paddle
import paddle.fluid as fluid

__all__ = [
    "PreActResNet"
]

from ppcls.modeling.architectures.cifar.layers import conv2d, pool2d, fc, bn_act, global_avg_pool, set_fp16


def shortcut(x, out_channels, stride, name):
    in_channels = x.shape[1]
    if in_channels != out_channels or stride != 1:
        if stride != 1:
            x = pool2d(x, 2, stride=2, type='avg', name=name + '.pool')
        x = conv2d(x, out_channels, 1, 1, name=name + '.conv')
    return x


def basic_block(x, out_channels, name):
    identity = x
    x = bn_act(x, 'relu', name + ".branch2a")
    x = conv2d(x, out_channels, kernel_size=3, name=name + ".branch2a")
    x = bn_act(x, 'relu', name + ".branch2b")
    x = conv2d(x, out_channels, kernel_size=3, name=name + ".branch2b")
    return fluid.layers.elementwise_add(x=identity, y=x)


def down_block(x, out_channels, stride, name):
    x = bn_act(x, 'relu', name + ".branch2a")
    identity = x
    x = conv2d(x, out_channels, kernel_size=3, stride=stride, name=name + ".branch2a")
    x = bn_act(x, 'relu', name + ".branch2b")
    x = conv2d(x, out_channels, kernel_size=3, name=name + ".branch2b")
    identity = shortcut(identity, out_channels, stride=stride, name=name + '.branch1')
    return fluid.layers.elementwise_add(x=identity, y=x)


class PreActResNet:

    def __init__(self, depth, k, fp16=False):
        self.depth = depth
        self.k = k
        set_fp16(fp16)

    def net(self, input, class_dim=10):
        depth = self.depth

        layers = [(depth - 4) // 6] * 3

        stages = [16, 32, 64]

        x = conv2d(input, 16, kernel_size=3)

        x = self._make_layer(x, stages[0] * self.k, layers[0], stride=1, name='stage1')
        x = self._make_layer(x, stages[1] * self.k, layers[1], stride=2, name='stage2')
        x = self._make_layer(x, stages[2] * self.k, layers[2], stride=2, name='stage3')

        x = bn_act(x, 'relu', name='post_activ')
        x = global_avg_pool(x)
        x = fc(x, class_dim)
        return x

    def _make_layer(self, x, out_channels, num_units, stride, name):
        x = down_block(x, out_channels, stride=stride, name=name + ".unit1")
        for i in range(1, num_units):
            x = basic_block(x, out_channels, name=name + f".unit{i+1}")
        return x